from langchain.prompts import PromptTemplate

word_list = [
    ["able", "ache", "ape", "bake", "blade", "blame", "brake"],
    ["brave", "cage", "cake", "cane", "case", "cave", "change"],
    ["chase", "date", "face", "fade", "fame", "flame", "frame"],
    ["game", "gate", "gaze", "grade", "grape", "haste", "hate"],
    ["lake", "lane", "late", "make", "male", "name", "pace"],
    ["page", "phase", "place", "plate", "race", "rate", "safe"],
    ["sale", "same", "save", "scale", "shade", "shame", "shape"],
    ["shave", "skate", "slave", "space", "spade", "stage", "stake"],
    ["strange", "take", "tale", "tame", "tape", "taste", "trade"],
    ["wage", "wake", "waste", "wave", "whale", "act", "add"],
    ["am", "and", "ash", "at", "axe", "bag", "band"],
    ["bank", "bat", "batch", "black", "brand", "can", "cap"],
    ["cash", "catch", "chat", "clap", "crab", "crash", "dad"],
    ["damn", "damp", "dash", "drag", "fact", "fan", "fat"],
    ["fax", "flag", "flash", "flat", "frank", "gang", "gap"],
    ["gas", "glad", "grab", "grand", "grant", "hack", "hang"],
    ["hat", "hatch", "jam", "jazz", "lab", "lack", "lad"],
    ["land", "lap", "mad", "man", "map", "mat", "match"],
    ["maths", "nap", "pan", "pant", "pants", "pat", "plan"],
    ["rack", "rag", "range", "rank", "rat", "sack", "sad"],
    ["salt", "sand", "slap", "slash", "smash", "snack", "snatch"],
    ["splash", "stack", "staff", "stamp", "strand", "swan"],
    ["tag", "tank", "tap", "tax", "than", "thank", "that"],
    ["track", "tram", "trap", "trash", "van", "wag", "want"],
    ["wash", "watch", "what", "wrap", "ask", "bath"],
    ["branch", "brass", "cast", "chance", "class", "craft"],
    ["dance", "draft", "fast", "glass", "grasp", "last"],
    ["pass", "past", "path", "path", "task"],
    ["aid", "aim", "aunt", "bay", "brain", "chain", "clause"],
    ["clay", "day", "drain", "fail", "faith", "fault", "gain"],
    ["gay", "grain", "gray", "jail", "laugh", "laungh"],
    ["mail", "main", "may", "nail", "paint", "pay"],
    ["pray", "rain", "ray", "sail", "saint", "sauce"],
    ["snail", "spray", "stain", "stay", "strain", "trail"],
    ["train", "tray", "waist", "way", "all"],
    ["ball", "calf", "call", "calm", "dawn", "drawer"],
    ["fall", "false", "hall", "jaw", "law", "lawn"],
    ["mall", "raw", "small", "stall", "tall", "walk"],
    ["wall", "yawn", "air", "arm", "art", "bar"],
    ["bare", "bark", "card", "care", "cart", "carve"],
    ["charge", "charm", "chart", "dare", "dark", "fair"],
    ["far", "fare", "farm", "hair", "hard", "harm"],
    ["jar", "large", "march", "mark", "marsh", "park"],
    ["part", "rare", "scare", "scarf", "share", "sharp"],
    ["spark", "stair", "star", "stare", "start"],
    ["starve", "ward", "warm", "yard"]
]

prompt_template = """
## 角色
你是一个英语老师，可以基于单词表来讲解单词。

## 单词表
{word_list}

## Commands - Prefix: "/"
- bm: 你的模型底座是什么
- cnt: 单词表中有多少个单词
- all: 显示<单词表>中的所有单词
- last: 显示<单词表>中的最后一个单词
"""
PROMPT = PromptTemplate(
    template=prompt_template, input_variables=["word_list"]
)

def get_word_list(row_cnt):
    prompt_word = ''
    for wlist in word_list[0:row_cnt]:
        word_line = ''
        for wl in wlist:
            word_line += wl.join(', ')
        prompt_word += word_line[1:] + '\n'
    return prompt_word

def get_word_prompt(row_cnt):
    prompt_word = get_word_list(row_cnt)
    vars = {'word_list': prompt_word}
    prompt_ret = PROMPT.format(**vars)
    return prompt_ret
